'''
Created on 9 janv. 2015
Date : 13 janv. 2015
Version : 29
@author: Durand Sebastien et Jomni Cyril
'''
# ==================== debut import
import numpy as np
import matplotlib.pyplot as plt
import time
from matplotlib.pyplot import winter, bone, pink, gray

# ==================== fin import

#fonction istencil :
def istencil_calcul(n, corel_range):
    #======================Generation d'un tableau de coordonnees Zx
    Zx = np.array(np.tile(np.arange(n),(n,1)), dtype = np.float64) + 1
    #print("Affichage de Zx\n")
    #print(Zx)
    
    #Generation du tableau de coordonnes Zy qui correspond, par definition, au transpose de Zx 
    Zy=Zx.T
    #print("Affichage de Zy\n")
    #print(Zy)
    
    #============================ debut calcul Xx et Xy 

    #vecteur 
    Xx_v = np.array(np.zeros(n)+n+1,dtype=np.float64)
    #print("Affichage de Xx_v\n")
    #print(Xx_v)
    
    Xy_v =np.array(np.arange(n)+1,dtype=np.float64)
    #print("Affichage de Xy_v\n")
    #print(Xy_v)
    
    # ============================== calcul de stencil
    m=long(np.log(n+1)/np.log(2) +1)
    
    stencil = np.zeros((n,n))
    stencil[0:n,0]=1
    for i in range(m):
        
        stencil[ ::corel_range**(i+1), i+1]=1
    
    stencil = np.roll(stencil,n/2, axis = 0)
    stencil = np.fliplr(stencil)
    istencil = np.where(stencil !=0)
    
    #print('affichage de istencil :')
    #print(istencil)
    
    return Zx,Zy,Xx_v,Xy_v,istencil



def covar(n,Zx,Zy,Xx_v,Xy_v,istencil):
    
    
    # calcul taille istencil
    taille_is = len(istencil[1])
    #============== calcul matrice Xx et Xy
    # version n/n
    Xx = np.tile(Xx_v,(n,1))
    #print('affichage de Xx :')
    #print(Xx)
    Xy = np.tile(Xy_v,(n,1))
    #print('affichage de Xy :')
    #print(Xy)
    
    # version n*2/n
    Xx2 = np.tile(Xx_v,(taille_is,1))
    #print('affichage de Xx2 :')
    #print(Xx2)
    Xy2 = np.tile(Xy_v,(taille_is,1))
    #print('affichage de Xy2 :')
    #print(Xy2)
    
    #print(Xx_v)
    #print(Xy_v)
    #print(taille_is)
    
    # =========================== calcul de Zxs  et Zys
    # creation des vecteur :
    Zxs_v = Zx[istencil] 
    Zys_v = Zy[istencil] 
        
    # creation des matrice n/n*2
    Zxsn = np.tile(Zxs_v,(n,1))
    #print('affichage de Zxsn :')
    #print(Zxsn)
    Zysn = np.tile(Zys_v,(n,1))
    #print('affichage de Zysn :')
    #print(Zysn)
    
    # creation des matrice n*2/n*2
    Zys = np.tile(Zys_v,(taille_is,1))
    #print('affichage de Zys :')
    #print(Zys)
    Zxs = np.tile(Zxs_v,(taille_is,1))
    #print('affichage de Zxs :')
    #print(Zxs)
    
    
    #  =========================  calcul de Zxref Zyref scalaire
    # derniere ligne et premiere colone
    Zxref = Zx[0,n-1]
    #print(Zxref)
    Zyref = Zy[0,n-1]
    #print(Zyref)
    
    #  ===================== calcul de Pxx
    
    Pxx = (-1*((Xx - Xx.T)**2+(Xy - Xy.T)**2)**(5./6.)) + \
        ((Zxref - Xx)**2 + (Zyref - Xy)**2)**(5./6.) + \
        (((Zxref - Xx)**2 + (Zyref - Xy)**2)**(5./6.)).T
        
    #print(Pxx.shape)
    
    # ==================== calcul de Pzz
    
    Pzz = (-1*((Zxs - Zxs.T)**2+(Zys - Zys.T)**2)**(5./6.)) + \
        ((Zxref - Zxs)**2 + (Zyref - Zys)**2)**(5./6.) + \
        (((Zxref - Zxs)**2 + (Zyref - Zys)**2)**(5./6.)).T
        
    #print(Pzz.shape)
    
    # ==================== calcul de Pxz 
    
    Pxz = (-1*((Xx2.T - Zxsn)**2+(Xy2.T - Zysn)**2)**(5./6.)) + \
        (((Zxref - Xx2)**2 + (Zyref - Xy2)**2)**(5./6.)).T + \
        ((Zxref - Zxsn)**2 + (Zyref - Zysn)**2)**(5./6.)
    
    #print(Pxz.shape)
    
    # ============== normalisation 
    
    Pxx = Pxx*0.5*6.88    
    Pzz = Pzz*0.5*6.88   
    Pxz = Pxz*0.5*6.88   
    
    #print('affichage de Pxx :')
    #print(Pxx)
    #print('affichage de Pzz :')
    #print(Pzz)
    #print('affichage de Pxz :')
    #print(Pxz)
    
    return Pxx, Pzz, Pxz

def calcul_A_B(n,corel_range):
    
    #Appelle des fonction istencil_calcul et covar :
    Zx,Zy,Xx_v,Xy_v,istencil = istencil_calcul(n,corel_range)
    Pxx, Pzz, Pxz = covar(n,Zx,Zy,Xx_v,Xy_v,istencil)
    
    #calcul talle de istencil :
    taille_is = len(istencil[1])
    
    # calcul SVD de Pzz car Pzz non inversible:
    U1,s,V1 = np.linalg.svd(Pzz)
    
    # filtrage de la derniere valeur propre
    s[taille_is-1]=1
    
    # inversion des valeurs propres :
    s = 1./s   
    s[taille_is-1]=0
    D = np.diag(s)
    
    # Calcul de Pzz inverse :
    Pzz_inv = np.dot(np.dot(V1.T,D),U1.T)

    # Calcul de A:
    A = np.dot(Pxz,Pzz_inv)

    #Calcul de B carre : Pxx - Pxz*Pzz_inv*Pzx:
    B_carre = Pxx - np.dot(A,Pxz.T)
    
    # SVD de B carre :
    U2,s2,V2 = np.linalg.svd(B_carre)
    
    # Calcul de la racine de la matrice diagonalise de B carre :
    D2 = np.diag(np.sqrt(s2))
    
    #Calcul de B :
    B = np.dot(U2,D2)
    
    return A, B, istencil

def iteration(A,B,X,n,dt):

    
    #Generation du vecteur G a partir d'un generateur de nombres aleatoires suivant une distribution normale, de variance unitaire et de moyenne nulle
    G = np.random.normal(0, 1, n)
    
    # definition de la valeur de reference
    # derniere ligne et premiere colone
    Zref = X[n-1,0]
    
    #Calcul de l'iteration
    
    #calcul de Z :
    Z = X[istencil]
    
    # Calcul de la colone d'apres
    ajout_col = (np.dot(A, (Z-Zref)) + np.dot(B,G)) + Zref
    
    
    #On ajoute la derniere colonne au tableau que l'on a sauvegarde
    for i in range(n-1):
        X[:,i] = X[:,i+1]
            
    X[:,n-1] = ajout_col

    #affichage
    #pause
    time.sleep(dt)
    #plot image
    plt.draw()
    # On fixe l'intervale pour accelerer l'affichage :
    plt.imshow(X,cmap=winter())
    #plt.imshow(X)
    
    return X


n = 8 # taille de l'image carre (n*n)
dt = 1/5. # temps de pause
n_it = 100  # nombre d'iteration
corel_range = 2 # taille de correlation
# plus il est petit et plus les cellules sont grandes

#generation de la matrice de depart
X =np.zeros((n,n))

#alumage du plot
plt.ion()

#Appelle de la fonction calcul_A_B pour A, B et istencil
A, B, istencil = calcul_A_B(n,corel_range)
#print('affichage de A :')
#print(A)

# iteration 
for i in range(n_it):
    # appelle la fonction d'iteration
    X = iteration(A,B,X,n,dt)
    
    
# fin